from typing import Dict, Any, Optional, Tuple
from datetime import timedelta
from utils.logger import Logger
from utils.exceptions import RateLimitError
from middleware.redis_adapter import get_redis_adapter
from utils.constants import DEFAULT_RATE_LIMIT  # 全局默认限流配置

# 初始化日志
logger = Logger.get_logger(agent_name="rate_limit_agent")


class RateLimitAgent:
    """限流Agent：支持接口级、Agent级、用户级限流"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.rate_limit_prefix = "rate_limit:"  # Redis键前缀
        # 默认限流配置（可通过Redis动态修改）
        self.default_limits = DEFAULT_RATE_LIMIT or {
            "api": {"limit": 100, "period": 60},  # 接口：60秒100次
            "agent": {"limit": 50, "period": 60},  # Agent实例：60秒50次
            "user": {"limit": 30, "period": 60}  # 用户：60秒30次
        }
        logger.info("限流Agent初始化成功，加载默认限流配置")

    def _get_rate_limit_key(self, limit_type: str, identifier: str) -> str:
        """获取限流Redis键（limit_type: 限流类型，identifier: 唯一标识）"""
        return f"{self.rate_limit_prefix}{limit_type}:{identifier}"

    def _get_limit_config(self, limit_type: str) -> Tuple[int, int]:
        """
        获取限流配置（优先Redis动态配置，无则用默认）
        Returns:
            (limit: 最大次数, period: 时间窗口（秒）)
        """
        try:
            # 从Redis获取动态配置
            config_key = f"{self.rate_limit_prefix}config:{limit_type}"
            config = self.redis_client.hgetall(config_key)
            if config:
                limit = int(config.get(b"limit", self.default_limits[limit_type]["limit"]))
                period = int(config.get(b"period", self.default_limits[limit_type]["period"]))
                return limit, period
            # 使用默认配置
            return self.default_limits[limit_type]["limit"], self.default_limits[limit_type]["period"]
        except Exception as e:
            logger.warning(f"获取限流配置失败，使用默认配置: limit_type={limit_type}, error={str(e)}")
            return self.default_limits[limit_type]["limit"], self.default_limits[limit_type]["period"]

    def check_rate_limit(self, limit_type: str, identifier: str) -> Dict[str, Any]:
        """
        检查限流（不抛出异常，返回限流结果）
        Args:
            limit_type: 限流类型（api/agent/user）
            identifier: 唯一标识（api路径/agent_id/user_id）
        Returns:
            {"allowed": bool, "remaining": int, "reset_after": int, "limit": int}
        """
        if limit_type not in self.default_limits:
            logger.warning(f"无效的限流类型: {limit_type}")
            return {"allowed": True, "remaining": -1, "reset_after": 0, "limit": -1}

        try:
            # 1. 获取限流配置
            limit, period = self._get_limit_config(limit_type)
            if limit <= 0:
                return {"allowed": True, "remaining": -1, "reset_after": 0, "limit": limit}

            # 2. 生成Redis键
            key = self._get_rate_limit_key(limit_type, identifier)

            # 3. Redis计数器+过期时间（原子操作）
            current_count = self.redis_client.incr(key)
            if current_count == 1:
                # 第一次计数，设置过期时间
                self.redis_client.expire(key, period)

            # 4. 计算剩余次数和重置时间
            remaining = max(0, limit - current_count)
            ttl = self.redis_client.ttl(key)  # 剩余过期时间（秒）
            reset_after = ttl if ttl > 0 else period

            # 5. 判断是否限流
            allowed = current_count <= limit
            if not allowed:
                logger.warning(
                    f"触发限流: type={limit_type}, id={identifier}, count={current_count}/{limit}, reset_after={reset_after}s")
            else:
                logger.debug(f"限流检查通过: type={limit_type}, id={identifier}, remaining={remaining}")

            return {
                "allowed": allowed,
                "remaining": remaining,
                "reset_after": reset_after,
                "limit": limit
            }
        except Exception as e:
            logger.error(f"限流检查失败: type={limit_type}, id={identifier}, error={str(e)}", exc_info=True)
            # 限流检查失败时默认允许请求（避免影响业务）
            return {"allowed": True, "remaining": -1, "reset_after": 0, "limit": -1}

    def enforce_rate_limit(self, limit_type: str, identifier: str) -> None:
        """
        强制限流（触发限流时抛出RateLimitError）
        Args:
            limit_type: 限流类型（api/agent/user）
            identifier: 唯一标识（api路径/agent_id/user_id）
        """
        result = self.check_rate_limit(limit_type, identifier)
        if not result["allowed"]:
            raise RateLimitError(
                message=f"触发{limit_type}级限流，{result['limit']}次/{result['reset_after']}秒，剩余{result['remaining']}次",
                context={"type": limit_type, "id": identifier, "limit": result["limit"],
                         "reset_after": result["reset_after"]}
            )

    def update_rate_limit_config(self, limit_type: str, limit: int, period: int) -> bool:
        """更新限流配置（动态修改，同步到Redis）"""
        if limit_type not in self.default_limits:
            logger.warning(f"无效的限流类型: {limit_type}")
            return False

        try:
            config_key = f"{self.rate_limit_prefix}config:{limit_type}"
            self.redis_client.hset(config_key, mapping={
                "limit": limit,
                "period": period
            })
            # 配置永久有效（除非手动删除）
            self.redis_client.persist(config_key)
            logger.info(f"更新限流配置成功: type={limit_type}, limit={limit}, period={period}")
            return True
        except Exception as e:
            logger.error(f"更新限流配置失败: type={limit_type}, error={str(e)}", exc_info=True)
            return False


# 单例实例
rate_limit_agent = RateLimitAgent()


# 简化外部调用函数
def check_rate_limit(limit_type: str, identifier: str) -> None:
    """限流检查快捷函数（触发限流抛出异常）"""
    rate_limit_agent.enforce_rate_limit(limit_type, identifier)


# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="rate_limit_test", agent_name="rate_limit_agent")
    test_logger.info("开始测试限流Agent")

    # 测试接口级限流（模拟60秒10次）
    rate_limit_agent.update_rate_limit_config("api", limit=10, period=60)
    test_api = "/agent/text_qa/predict"

    for i in range(12):
        try:
            check_rate_limit("api", test_api)
            test_logger.info(f"第{i + 1}次请求: 限流检查通过")
        except RateLimitError as e:
            test_logger.error(f"第{i + 1}次请求: {e}")

    # 测试Agent级限流
    test_agent_id = "text_qa_agent_001"
    rate_limit_agent.update_rate_limit_config("agent", limit=5, period=60)

    for i in range(7):
        try:
            check_rate_limit("agent", test_agent_id)
            test_logger.info(f"Agent请求第{i + 1}次: 限流检查通过")
        except RateLimitError as e:
            test_logger.error(f"Agent请求第{i + 1}次: {e}")

    test_logger.success("限流Agent测试完成")